# 只保留文件夹批量转换功能
# 用法:
#   python hdf52zarr.py <hdf5_folder> [--to_path <target.zarr>]
# 参数说明:
# - from_path: 源 HDF5 文件夹路径。
# - --to_path: 目标 Zarr 文件路径（可选，默认与源路径同名加后缀 -demo.zarr）。
# 示例:
#   python hdf52zarr.py data/hdf5_folder --to_path output/result.zarr
# 注意事项:
# - 支持递归复制 HDF5 的分组、数据集及其属性到 Zarr。
# - 若目标 Zarr 路径所在目录不存在会自动创建。
# - 所有 HDF5 文件内容会合并到同一个 Zarr 文件中。
import argparse
import h5py
import zarr
import numpy as np
import os
from pathlib import Path
from scipy.spatial.transform import Rotation

def rechunk(zarr_path:str, chunk_size:int =100):
    def _rechunk_group(group):
        for key, item in group.items():
            if isinstance(item, zarr.Group):
                _rechunk_group(item)
            elif isinstance(item, zarr.Array):
                arr = item
                if isinstance(arr.chunks, tuple):
                    target_chunks = tuple([chunk_size] + list(arr.chunks[1:]))
                else:
                    target_chunks = (chunk_size,)
                if arr.chunks != target_chunks:
                    tmp_name = key + "_tmp_rechunk"
                    if tmp_name in group:
                        del group[tmp_name]
                    tmp_arr = group.create_dataset(tmp_name, shape=arr.shape, dtype=arr.dtype, chunks=target_chunks, overwrite=True)
                    tmp_arr[:] = arr[:]
                    del group[key]
                    group.move(tmp_name, key)
    root = zarr.open_group(zarr_path, mode='a')
    _rechunk_group(root)

def hdf5_to_dict(h5_file):
    """递归读取h5文件内容为dict"""
    def read_group(g):
        out = {}
        for k, v in g.items():
            if isinstance(v, h5py.Dataset):
                out[k] = v[()]
            elif isinstance(v, h5py.Group):
                out[k] = read_group(v)
        return out
    return read_group(h5_file)

def hdf5_to_zarr(hdf5_path, zarr_path, mode, episode_offset=0):
    with h5py.File(hdf5_path, 'r') as h5_file:
        data = hdf5_to_dict(h5_file)
    root = zarr.open_group(zarr_path, mode=mode)
    # 确保分组存在
    if mode == 'w':
        root.create_group('data')
        root.create_group('pcd')
        root.create_group('ee')
        root.create_group('r_pcd_output')
    # meta/episode_ends
    episode_len = data['action'].shape[0]
    meta_group = root.require_group('meta')
    if mode == 'w':
        meta_group.array('episode_ends', [episode_len])
    else:
        episode_ends = meta_group['episode_ends']
        episode_ends.append([episode_len + episode_ends[-1]])
    # qpos/action
    qpos_group = root.require_group('qpos')
    action = data['action']
    qpos_state = data['observations']['qpos']
    if mode == 'a':
        qpos_group['action'].append(action)
        qpos_group['state'].append(qpos_state)
    else:
        qpos_group.array('action', action)
        qpos_group.array('state', qpos_state)
    # cameras
    cameras_group = root.require_group('cameras')
    # midBack_camera/rgb
    midBack_camera = cameras_group.require_group('midBack_camera')
    cam_high = data['observations']['images']['cam_high']
    if mode == 'a':
        midBack_camera['rgb'].append(cam_high)
    else:
        midBack_camera.array('rgb', cam_high)
    # midLeft_camera/rgb
    midLeft_camera = cameras_group.require_group('midLeft_camera')
    cam_left = data['observations']['images']['cam_left_wrist']
    if mode == 'a':
        midLeft_camera['rgb'].append(cam_left)
    else:
        midLeft_camera.array('rgb', cam_left)
    # midRight_camera/rgb
    midRight_camera = cameras_group.require_group('midRight_camera')
    cam_right = data['observations']['images']['cam_right_wrist']
    if mode == 'a':
        midRight_camera['rgb'].append(cam_right)
    else:
        midRight_camera.array('rgb', cam_right)
    print(f"Converted {hdf5_path} to {zarr_path}")

def hdf5_folder_to_zarr(hdf5_folder, zarr_path):
    flag = False
    for filename in sorted(os.listdir(hdf5_folder)):
        if filename.endswith('.hdf5') or filename.endswith('.h5'):
            hdf5_path = os.path.join(hdf5_folder, filename)
            if not flag:
                hdf5_to_zarr(hdf5_path, zarr_path, mode='w')
                flag = True
            else:
                hdf5_to_zarr(hdf5_path, zarr_path, mode='a')
    rechunk(zarr_path)
    print(f"All HDF5 files in {hdf5_folder} converted to {zarr_path}")

def main():
    parser = argparse.ArgumentParser(description="Convert folder of HDF5 files to Zarr")
    parser.add_argument("from_path", type=str, help="Source HDF5 folder path")
    parser.add_argument("--to_path", type=str, help="Target Zarr file path")
    args = parser.parse_args()
    if args.to_path is None:
        args.to_path = str(Path(args.from_path).parent) + '/' + str(Path(args.from_path).stem) + '-demo.zarr'
    zarr_dir = os.path.dirname(args.to_path)
    if zarr_dir and not os.path.exists(zarr_dir):
        os.makedirs(zarr_dir, exist_ok=True)
    hdf5_folder_to_zarr(args.from_path, args.to_path)

if __name__ == "__main__":
    main()